package state

import (
	"fmt"
	"osiris/dto"
	"osiris/errors"
	"osiris/logger"
	"osiris/ocrypto/ecdsa"
	"sync"

	peer "github.com/libp2p/go-libp2p/core/peer"
)

// var 状态树（只有当收到区块或创造区块时才会涉及到状态树的写操作）
var (
	accountTrie *AccountStateTrie
	peerTrie    *PeerStateTrie
)

// InitState 初始化所有状态树
func InitState() {
	accountTrie = &AccountStateTrie{
		BaseStateTrie: BaseStateTrie{
			ChainID: 0,
		},
	}

	peerTrie = &PeerStateTrie{
		BaseStateTrie: BaseStateTrie{
			ChainID: 0,
		},
	}

	accountTrie.InitState()
	peerTrie.InitState()
}

// ResetState 重置所有状态树
func ResetState() {
	accountTrie.ResetState()
	peerTrie.ResetState()
}

// buildTxHandler 根据txData.TxType构建对应的TxHandler，方便后续的Verify, Commit, Undo
//
//	@param txData
//	@return TxHandler 可能为nil，记得判空
func buildTxHandler(txData dto.TxData) TxHandler {
	switch txData.TxType {
	case dto.TX_REGISTER:
		return RegisterTxHandler{
			TxData: txData,
		}

	case dto.TX_DEPOSIT:
		return DepositTxHandler{
			TxData: txData,
		}

	case dto.TX_TRANSFORM:
		return TransformTxHandler{
			TxData: txData,
		}

	case dto.TX_PLEDGE:
		return PledgeTxHandler{
			TxData: txData,
		}

	case dto.TX_RELEASE:
		return ReleaseTxHandler{
			TxData: txData,
		}

	case dto.TX_REWARD:
		return RewardTxHandler{
			TxData: txData,
		}

	default:
		logger.Warn(map[string]interface{}{"[state] [build TxHandler]": "unknown Tx type"})
		return nil
	}
}

// getAccountTrie 根据交易信息将交易分流到对应的账户状态账本（状态树）
//
//	@param chainID 属于哪条链
//	@return *AccountStateTrie
func getAccountTrie(chainID int) *AccountStateTrie {
	//TODO 多通道，chainID
	return accountTrie
}

// getPeerTrie 根据交易信息将交易分流到对应的节点状态账本（状态树）
//
//	@param chainID 属于哪条链
//	@return *PeerStateTrie
func getPeerTrie(chainID int) *PeerStateTrie {
	//TODO 多通道，chainID
	return peerTrie
}

// setAccountTrie 修改对应chainID的账户状态树指针
//
//	@param chainID
//	@param tempTrie
func setAccountTrie(chainID int, tempTrie *AccountStateTrie) {
	//TODO 多通道，chainID
	accountTrie.Lock()
	accountTrie = tempTrie
	accountTrie.UnLock()
}

// setPeerTrie 修改对应chainID的节点状态树的指针
//
//	@param chainID
//	@param tempTrie
func setPeerTrie(chainID int, tempTrie *PeerStateTrie) {
	//TODO 多通道，chainID
	peerTrie.Lock()
	peerTrie = tempTrie
	peerTrie.UnLock()
}

// GetAccountRollbackChan 建立一个账户状态树的备份，并返回一个通道用于接收是否回滚状态树
//
//	@param chainID
//	@return chan 如果chan传来true，则回滚到当初建立的状态树备份
func GetAccountRollbackChan(chainID int) chan bool {
	rollbackChan := make(chan bool)
	wg := new(sync.WaitGroup)
	wg.Add(1)

	go func(tempChainID int) {
		tempTrie := getAccountTrie(tempChainID)
		tempTrie.RLock()
		tempTrieVal := *tempTrie
		tempTrieCopy := &tempTrieVal
		accountTrie.RUnLock()
		wg.Done()

		needRollback := <-rollbackChan
		if needRollback {
			setAccountTrie(tempChainID, tempTrieCopy)
		}
	}(chainID)

	wg.Wait()
	return rollbackChan
}

// GetPeerRollbackChan 建立一个节点状态树的备份，并返回一个通道用于接收是否回滚状态树
//
//	@param chainID
//	@return chan 如果chan传来true，则回滚到当初建立的状态树备份
func GetPeerRollbackChan(chainID int) chan bool {
	rollbackChan := make(chan bool)
	wg := new(sync.WaitGroup)
	wg.Add(1)

	go func(tempChainID int) {
		tempTrie := getPeerTrie(tempChainID)
		tempTrie.RLock()
		tempTrieVal := *tempTrie
		tempTrieCopy := &tempTrieVal
		tempTrie.RUnLock()
		wg.Done()

		needRollback := <-rollbackChan
		if needRollback {
			setPeerTrie(tempChainID, tempTrieCopy)
		}
	}(chainID)

	wg.Wait()
	return rollbackChan
}

// verifyTx 验证交易内容的状态转换是否合法（不包括密码学类的验证，密码学类的验证在dto包中完成；只会读状态树，不需要加错误回滚）
//
//	@param txData
//	@return bool 验证是否通过
func verifyTx(txData dto.TxData) bool {
	txHandler := buildTxHandler(txData)
	if txHandler == nil {
		return false
	}

	trie1 := getAccountTrie(txData.ChainID)
	trie2 := getPeerTrie(txData.ChainID)

	return txHandler.Verify(trie1, trie2)
}

// commitTx 根据交易改变状态
//
//	@param txData
//	@return error
func commitTx(txData dto.TxData) error {
	txHandler := buildTxHandler(txData)
	if txHandler == nil {
		return &errors.UnknownTxTypeError{}
	}

	trie1 := getAccountTrie(txData.ChainID)
	trie2 := getPeerTrie(txData.ChainID)

	success := txHandler.Commit(trie1, trie2)
	if success {
		return nil
	} else {
		return &errors.TxCommitError{}
	}
}

// CommitTxBatch 用于记账：尝试批量提交交易（单笔单笔地提交交易，某笔交易发生错误时，会回滚该笔交易执行前的状态）
//
//	@param txs
//	@return []dto.TxData 执行成功的交易
func CommitTxBatch(txs []dto.TxData) []dto.TxData {
	succeessTxs := make([]dto.TxData, 0)

	for _, tx := range txs {
		valid := verifyTx(tx)
		if !valid {
			continue
		}

		err := commitTx(tx)
		if err != nil {
			logger.Error(map[string]interface{}{fmt.Sprintf("[state] [Commit Tx Batch] %s", tx.TxHash): err})
			continue
		}
		succeessTxs = append(succeessTxs, tx)
	}

	return succeessTxs
}

// TryCommitTxBatch 用于执行新区块中的交易（但凡有一笔交易发生错误，会回滚至所有交易执行之前的状态）
//
//	@param chainID
//	@param txs
//	@return error 碰到的第一个错误
func TryCommitTxBatch(chainID int, txs []dto.TxData) error {
	accountTrieVal := *(getAccountTrie(chainID))
	accoutTrieCopy := &accountTrieVal
	peerTrieVal := *(getPeerTrie(chainID))
	peerTrieCopy := &peerTrieVal

	errOccured := false
	defer func() {
		if errOccured {
			setAccountTrie(chainID, accoutTrieCopy)
			setPeerTrie(chainID, peerTrieCopy)
		}
	}()

	for _, tx := range txs {
		valid := verifyTx(tx)
		if !valid {
			continue
		}

		err := commitTx(tx)
		if err != nil {
			logger.Error(map[string]interface{}{fmt.Sprintf("[state] [Commit Tx Batch] %s", tx.TxHash): err})
			errOccured = true
			return err
		}
	}

	return nil
}

// GetWriterPeer 获得记账节点
//
//	@param chainID
//	@return peer.ID
//	@return error
func GetWriterPeer(chainID int) (peer.ID, error) {
	tempPeerTrie := getPeerTrie(chainID)
	if tempPeerTrie == nil {
		err := &errors.UnknownPeerTrieError{}
		logger.Error(map[string]interface{}{"[state] [Get Writer Peer] ": err})
		return "", err
	}

	return tempPeerTrie.getHighestStake()
}

func CoolDownPeer(chainID int, peerIDBytes []byte) error {
	tempPeerTrie := getPeerTrie(chainID)
	if tempPeerTrie == nil {
		err := &errors.UnknownPeerTrieError{}
		logger.Error(map[string]interface{}{"[state] [Cool Down Peer] ": err})
		return err
	}

	return tempPeerTrie.coolDownPeer(peerIDBytes)
}

func WarmUpPeers(chainID int) error {
	tempPeerTrie := getPeerTrie(chainID)
	if tempPeerTrie == nil {
		err := &errors.UnknownPeerTrieError{}
		logger.Error(map[string]interface{}{"[state] [Cool Down Peer] ": err})
		return err
	}

	return tempPeerTrie.warmUpPeers()
}

// UndoTx 根据交易回滚状态
//
//	@param txData
//	@return error
func UndoTx(txData dto.TxData) error {
	txHandler := buildTxHandler(txData)
	if txHandler == nil {
		return &errors.UnknownTxTypeError{}
	}

	trie1 := getAccountTrie(txData.ChainID)
	trie2 := getPeerTrie(txData.ChainID)

	success := txHandler.Undo(trie1, trie2)
	if success {
		return nil
	} else {
		return &errors.TxUndoError{}
	}
}

// GetAccountNonce 得到账户当前状态的nonce
//
//	@param addrStr 账户地址
//	@param chainID
//	@param txType
//	@return int64
func GetAccountNonce(addrStr string, chainID int) int64 {
	trie := getAccountTrie(chainID)
	if trie == nil {
		return InvalidNonce
	}

	nonce, exist := trie.getNonce(addrStr)
	if !exist {
		return InvalidNonce
	}

	return nonce
}

func HasAccountLeaf(addrBytes []byte, chainID int) bool {
	trie := getAccountTrie(chainID)
	if trie == nil {
		return false
	}

	return trie.HasLeaf(addrBytes)
}

func GetAccountLeaf(addrBytes []byte, chainID int) (dto.StateLeaf, bool) {
	trie := getAccountTrie(chainID)
	if trie == nil {
		return dto.StateLeaf{}, false
	}

	return trie.GetLeaf(addrBytes)
}

func HasPeerLeaf(peerIDBytes []byte, chainID int) bool {
	trie := getPeerTrie(chainID)
	if trie == nil {
		return false
	}

	return trie.HasLeaf(peerIDBytes)
}

func GetPeerLeaf(peerIDBytes []byte, chainID int) (dto.StateLeaf, bool) {
	trie := getPeerTrie(chainID)
	if trie == nil {
		return dto.StateLeaf{}, false
	}

	return trie.GetLeaf(peerIDBytes)
}

func AddPeerLeaf(chainID int, peerIDBytes []byte) error {
	trie := getPeerTrie(chainID)
	if trie == nil {
		err := &errors.UnknownPeerTrieError{}
		logger.Error(map[string]interface{}{"[state] [Add Peer Leaf]": err})
		return err
	}

	if trie.HasLeaf(peerIDBytes) {
		err := &errors.PeerExistError{}
		logger.Error(map[string]interface{}{"[p2p] [Add Peer Leaf]": err})
		return err
	}

	leaf := dto.StateLeaf{
		Asset:     0,
		CoolDown:  false,
		CoinSince: 0,
	}

	trie.Lock()
	trie.commitState(peerIDBytes, leaf, false)
	trie.UnLock()

	return nil
}

// GetAccountTrieRoot 账户状态树的根hash
//
//	@param chainID
//	@return string
//	@return error
func GetAccountTrieRoot(chainID int) (string, error) {
	tempAccountTrie := getAccountTrie(chainID)
	if tempAccountTrie == nil {
		err := &errors.UnknownAccountTrieError{}
		logger.Error(map[string]interface{}{"[state] [Get Account Trie Root] get account trie": err})
		return "", err
	}

	hashBytes := tempAccountTrie.RootHash()
	return ecdsa.Keccak256HashBytesToStr([]byte(hashBytes)), nil
}

// GetPeerTrieRoot 节点状态树的根hash
//
//	@param chainID
//	@return string
//	@return error
func GetPeerTrieRoot(chainID int) (string, error) {
	tempPeerTrie := getPeerTrie(chainID)
	if tempPeerTrie == nil {
		err := &errors.UnknownPeerTrieError{}
		logger.Error(map[string]interface{}{"[state] [Get Peer Trie Root] get peer trie": err})
		return "", err
	}

	hashBytes := tempPeerTrie.RootHash()
	return ecdsa.Keccak256HashBytesToStr([]byte(hashBytes)), nil
}

func IsTxNonceValid(chainID int, txType dto.TxType, addrStr string, nonce int64) TxStatus {
	if txType == dto.TX_RELEASE || txType == dto.TX_REGISTER {
		return TxPending
	}

	trie := getAccountTrie(chainID)
	if trie == nil {
		return TxInvalid
	}

	return trie.IsNonceValid(addrStr, nonce)
}
